# 导入必要的库和模块
import gradio as gr
import numpy as np
import pickle
from PIL import Image
from pinecone import Pinecone

# 初始化 Pinecone
pinecone_api_key = "e4df12f5-ebc2-4814-b493-3e166d6de302"  # 替换为您的Pinecone API密钥
pinecone_index_name = "mnist-index"  # 替换为您的索引名称

# 连接到 Pinecone
pinecone = Pinecone(api_key=pinecone_api_key)
index = pinecone.Index(pinecone_index_name)


# 定义预测函数，这个函数将用于Gradio接口进行预测
def predict_digit(image_data):
    img_array = np.array(image_data, dtype=np.uint8)
    # 将输入图像转换为灰度并调整大小
    img = Image.fromarray(img_array).convert('L').resize((28, 28))

    # 将图像转换为numpy数组并进行归一化
    img_array = np.array(img) // 16  # 缩放到0-15
    img_array = img_array.flatten().tolist()  # 扁平化并转换为列表

    # 使用准备好的查询向量在 Pinecone 索引中执行搜索
    results = index.query(vector=img_array, top_k=11, include_metadata=True)

    # 从搜索结果中提取每个匹配项的标签
    labels = [match['metadata']['label'] for match in results['matches']]

    # 使用投票机制确定最终的分类结果
    if labels:
        final_prediction = max(set(labels), key=labels.count)  # 使用最频繁的标签作为预测结果
        return int(final_prediction)
    else:
        return "未找到匹配项"


# 创建Gradio接口，这个接口将用于用户输入和显示预测结果
interface = gr.Interface(
    fn=predict_digit,  # 传入预测函数
    inputs=gr.Sketchpad(),  # 输入为用户书写的图像
    outputs=gr.Label(),  # 输出为预测结果
    title="手写数字识别",
    description="用鼠标书写数字，系统将预测数字。"
)

# 启动Gradio接口，用户可以通过这个接口进行交互
interface.launch()
